import torch
from torch.utils.data import Dataset
import os
import numpy as np


MAX_FREQ = 25


class ModalDataset(Dataset):
    def __init__(self, modaldata_dir, train_test='train'):
        self.filesnames = get_files_by_extension(modaldata_dir, '.npz')
        self.modaldata_dir = modaldata_dir
        self.train_test = train_test

    def __len__(self):
        if self.train_test == 'train':
            return len(self.filesnames[::3])
        elif self.train_test == 'test':
            return len(self.filesnames[1::6])

    def __getitem__(self, idx):
        if self.train_test == 'train':
            idx = idx * 3
        elif self.train_test == 'test':
            idx = idx * 6
        filename = os.path.join(self.modaldata_dir, self.filesnames[idx])
        data = np.load(filename)
        nodegrid = torch.from_numpy(data['nodegrid']).float()
        freqs = torch.from_numpy(data['freqs']).float()
        freqs = freqs.unsqueeze(0).expand(nodegrid.shape[0],  -1) / MAX_FREQ
        cellstate = torch.from_numpy(data['cellgrid']).long()
        total_modal_num = (nodegrid.shape[2] - 4) // 3
        modalinfo = nodegrid[:, :, -total_modal_num:]
        modalinfo = torch.flip(modalinfo, dims=[0])
        for mode in range(total_modal_num):
            modalinfo[:, :, mode] = modalinfo[:, :, mode] / torch.max(torch.abs(modalinfo[:, :, mode]))
  
        modalinfo = modalinfo.reshape(modalinfo.shape[0], -1)
      
        modalinfo = torch.concat((freqs, modalinfo), axis=1)

        return modalinfo, cellstate
    

def get_files_by_extension(folder_path, extension):
    """获取文件夹中指定扩展名的所有文件"""
    files = []
    for file in os.listdir(folder_path):
        if file.endswith(extension):
            files.append(file)
    return files